#' 
#' @TODO R包绘制森林图
#' @title ## R包绘制森林图
#' @param clinical 临床信息文件
#' @param saveplot 是否在输出目录保存图片
#' 
#' @examples 
#'  表达谱矩阵行名需要与生存信息第一列保持一致，要么都是带有01A的，保证没有重复的样本信息，要么都是没有01A的个人ID，如TCGA-NJ-A4YI
#'
#' > HNSC_clinical
#'               sample status time riskgroup riskscore age gender        stage
#' 1   TCGA-HD-7831-01A      0  667      High 1.6416701  74   male Stage III&IV
#' 2   TCGA-H7-7774-01A      0  407       Low 1.1163304  75 female Stage III&IV
#' 3   TCGA-CV-6943-01A      1  602       Low 0.7322320  74   male Stage III&IV
#' 4   TCGA-CN-5374-01A      1 1732       Low 1.2516243  56 female Stage III&IV
#' 5   TCGA-CV-6959-01A      1  256      High 1.8382042  48   male Stage III&IV
#'
#' @param out_dir 结果输出目录，生成figure与table文件夹，分别存储图片与表格文件,字符串向量，必须以 **"/"** 结尾
#' @param var_name 用于命名文件夹，如果为NULL，则使用clinical变量名
#' @param width 图片保存宽度
#' @param points_color 森林图中间points颜色，默认"#005670e3"
#' @param multi_after_univcox 是否使用单因素中显著的因素进行多因素分析，默认FALSE，如果为T，那么只支持二分类变量，不支持多个分类的单因素分析。
#' @param features_col 字符串向量,纳入分析的特征名称，必须为`clinical`中的列名
#' @usage 
#' a <- univ_multi_Cox_forestplot(
#'  clinical = aa, saveplot = T, out_dir = output_dir,
#'  features_col = c("riskgroup", "Age", "Gender", "Stage", "grade", "smoking", "margin_status", "alcohol_history", "hpv_status"),
#'  var_name = "HNSC", width = 7.5, multi_after_univcox = T
#' )
#' 
#' @return  *list*
#' @author: WYK
#' 
univ_multi_Cox_forestplot <- function(clinical = NULL, saveplot = F, out_dir = "./",features_col = NULL,
                                      var_name = NULL, width = 7.5, points_color = "#1d665a",multi_after_univcox = F,
                                      pval = 0.05) {
  if (is.null(var_name)) {
    var_name <- sample(letters,4) %>% paste0(.,collapse = '')
  }

  univCox_forestplot_res <- univCox_forestplot(clinical = clinical, points_color = points_color,features_col = features_col)

  if (multi_after_univcox) {

    features_col_multicox_index <- which(univCox_forestplot_res[["univcox_clinical_res"]][['p.value']] < pval) 

    if (length(features_col_multicox_index) >= 2) {
      multiCox_forestplot_res <- multiCox_forestplot(
        clinical = clinical,
        points_color = points_color,
        features_col = features_col[features_col_multicox_index]
      )
    } else if (length(features_col_multicox_index) <= 1) {
      multiCox_forestplot_res <- multiCox_forestplot(
        clinical = clinical,
        points_color = points_color,
        features_col = features_col
      )
    } else {
      multiCox_forestplot_res <- multiCox_forestplot(
        clinical = clinical,
        points_color = points_color,
        features_col = features_col
      )
    }

  } else {
    multiCox_forestplot_res <- multiCox_forestplot(clinical = clinical, points_color = points_color,features_col = features_col)
  }

#  + theme(plot.margin = unit(c(0.1, 0.1, 0.1, 0.1), "cm"))
#   + theme(plot.margin = unit(c(-.15, 0.1, 0.1, 0.1), "cm"))
  p <- cowplot::plot_grid(univCox_forestplot_res$plot,
    multiCox_forestplot_res$plot,nrow = 2, align = "v", scale = c(.98, .98),
    rel_heights = c(nrow(univCox_forestplot_res[[1]]), nrow(multiCox_forestplot_res[[1]]))
  )
  # print(p)

  res_all <- append(univCox_forestplot_res, multiCox_forestplot_res)

  if (saveplot) {
    if (!dir.exists(sprintf("%s/forestplot/%s", out_dir, var_name))) {
      dir.create(sprintf("%s/forestplot/%s", out_dir, var_name), showWarnings = F, recursive = T)
    } else {
      print(sprintf("Dir '%s/forestplot/%s' is existed.", out_dir, var_name))
    }
    dir_now <- sprintf("%s/forestplot/%s", out_dir, var_name)

    write_delim(
      x = univCox_forestplot_res[[1]],
      file = sprintf("%s/forestplot/%s/Table_Univcox_%s.txt", out_dir, var_name, var_name),
      delim = "\t", quote = "none"
    )
    write_delim(
      x = multiCox_forestplot_res[[1]],
      file = sprintf("%s/forestplot/%s/Table_Multicox_%s.txt", out_dir, var_name, var_name),
      delim = "\t", quote = "none"
    )

    plotout(
      od = sprintf("%s/forestplot/%s", out_dir, var_name),
      w = width, h = (nrow(univCox_forestplot_res[[1]]) * .73),
      p = univCox_forestplot_res$plot, name = "Univcox"
    )

    # cowplot::ggsave2(
    #   file = sprintf("%s/forestplot/Figure_univ_multi_cox_%s/Figure_Univcox_%s.pdf", out_dir, var_name, var_name),
    #   plot = univCox_forestplot_res$plot,
    #   width = width, height = (nrow(univCox_forestplot_res[[1]]) * .73)
    # )
    plotout(
      od = sprintf("%s/forestplot/%s", out_dir, var_name),
      w = width, h = (nrow(multiCox_forestplot_res[[1]]) * .73),
      p = multiCox_forestplot_res$plot, name = "MultiCox"
    )
    # cowplot::ggsave2(
    #   file = sprintf("%s/forestplot/Figure_univ_multi_cox_%s/Figure_MultiCox_%s.pdf", out_dir, var_name, var_name),
    #   plot = multiCox_forestplot_res$plot,
    #   width = width, height = (nrow(multiCox_forestplot_res[[1]]) * .73)
    # )

    plotout(
      od = sprintf("%s/forestplot/%s", out_dir, var_name),
      w = width, h = (nrow(univCox_forestplot_res[[1]])  *.7 + nrow(multiCox_forestplot_res[[1]]) * .7),
      p = p, name = "Univ_MultiCox"
    )

    # cowplot::ggsave2(
    #   filename = sprintf("%s/forestplot/Figure_univ_multi_cox_%s/Fig_%s.pdf", out_dir, var_name, var_name),
    #   plot = p, width = width, height = (nrow(univCox_forestplot_res[[1]])  *.7 + nrow(multiCox_forestplot_res[[1]]) * .7)
    # )
    # cowplot::ggsave2(
    #   filename = sprintf("%s/forestplot/Figure_univ_multi_cox_%s/Fig_%s.tiff", out_dir, var_name, var_name),
    #   plot = p, width = width, height = (nrow(univCox_forestplot_res[[1]])  *.7 + nrow(multiCox_forestplot_res[[1]]) * .7), dpi = 300
    # )
    # cowplot::ggsave2(
    #   filename = sprintf("%s/forestplot/Figure_univ_multi_cox_%s/Fig_%s_72dpi.tiff", out_dir, var_name, var_name),
    #   plot = p, width = width, height = (nrow(univCox_forestplot_res[[1]])  *.7 + nrow(multiCox_forestplot_res[[1]]) * .7), dpi = 72
    # )
  } else {
    print("Save plot by Parameter 'saveplot = T'.")
  }

  names(res_all) <- c('Table_Univcox','Figure_Univcox','Table_Multicox','Multi_Coxph_res','Figure_MultiCox')

  return(res_all)
}
# 子函数，单因素COX回归分析，临床特征间的比较
univCox_forestplot <- function(clinical = NULL, features_col = NULL,points_color = "#005670e3") {
  require(forestmodel)
  require(tidyverse)
  require(survival)
  require(survminer)

  vars_for_table <- clinical %>%
    dplyr::select(any_of(features_col)) %>%
    colnames(.)

  univ_formulas <- sapply(vars_for_table, function(x) as.formula(paste("Surv(time, status)~", x)))

  univ_models <- lapply(univ_formulas, function(x)  coxph(x, data = clinical))

  univcox_clinical_res <- map_dfr(.x = univ_models, .f = function(x) {
    result <- summary(x)

    result2 <- data.frame(
      var = rownames(result$coefficients),
      coef = result$coefficients[, 1],
      Hazard_Ratio = result$coefficients[, 2],
      se = result$coefficients[, 3],
      # p.value = result$coefficients[, 5],
      p.value = result$sctest["pvalue"],
      lower_.95 = result$conf.int[, "lower .95"],
      upper_.95 = result$conf.int[, "upper .95"],
      logrank_pvalue = result$sctest["pvalue"],
      wald_pvalue = result$waldtest["pvalue"],
      Likelihood_pvalue = result$logtest["pvalue"]
    )

    result2 <- result2 %>%
      mutate(HR = paste0(
        sprintf("%.2f", Hazard_Ratio),
        "(",
        sprintf("%.2f", lower_.95),
        "-",
        sprintf("%.2f", upper_.95),
        ")"
      ))
  })

  str(univcox_clinical_res)
  rownames(univcox_clinical_res) <- univcox_clinical_res$var
  
  panels <- list(
      list(width = 0.01),
      list(width = 0.1, display = ~ variable, fontface = "bold", heading = "Variable"),
      list(width = 0.02),
      list(width = 0.15, display = ~ level),
      list(width = 0.1, display = ~n, hjust = 1, heading = "N"),
      # list(width = 0.05, display = ~n_events, width = 0.05, hjust = 1, heading = "Events"),
      # list(
      #   width = 0.05,
      #   display = ~ replace(sprintf("%0.1f", person_time / 365.25), is.na(person_time), ""),
      #   heading = "Person-\nYears", hjust = 1
      # ),
      # format.pval(p.value, digits = 2L, eps = 0.001))
      list(width = 0.025, item = "vline", hjust = 0.5),
      list(
        width = 0.25, item = "forest", hjust = 0.5, heading = "Univariable Cox regression", linetype = "dashed",
        line_x = 0, fontsize = 3.73
      ),
      list(width = 0.025, item = "vline", hjust = 0.5),
      list(width = 0.15, display = ~ ifelse(reference, "Reference", sprintf(
        "%0.2f (%0.2f - %0.2f)",
        trans(estimate), trans(conf.low), trans(conf.high)
      )), display_na = NA, hjust = 0, heading = "HR(95% CI)"),
      list(width = 0.04),
      list(
        width = 0.2,
        display = ~ ifelse(reference, "",case_when(p.value < 0.001 ~ "<0.001",
        p.value > 0.001 ~ sprintf("%.3f",p.value))),
        display_na = NA, hjust = 1, heading = "p"
      ),
      list(width = 0.01)
  )

  format_options <- forest_model_format_options(
    colour = points_color,
    color = NULL,
    shape = 18,
    text_size = 3.73,
    point_size = 4.25,
    banded = TRUE
  )

  p <- forest_model(
    model_list = univ_models,
    panels = panels,
    # covariates = vars_for_table,
    merge_models = T,
    return_data = T,
    recalculate_height = T,
    recalculate_width = T,
    format_options = format_options,
    theme = theme_forest()
  )

  # print(p$plot)

  tmp <- list()
  tmp[[1]] <- univcox_clinical_res
  tmp[[2]] <- p$plot

  names(tmp) <- c("univcox_clinical_res", "plot")
  return(tmp)
}
# 子函数，多因素COX回归分析，临床特征间的比较
multiCox_forestplot <- function(clinical = NULL, features_col = NULL,points_color = "#005670e3") {
  require(forestmodel)
  require(tidyverse)
  require(survival)
  require(survminer)
  require(cowplot)

  vars_for_table <- clinical %>%
    dplyr::select(any_of(features_col)) %>%
    colnames(.)

  multicox_formulas <-
    as.formula(paste(
      "Surv(time, status)~",
      paste0(sep = "`", vars_for_table, sep = "`", collapse = "+")
    ))
  # top 4 colunm：sample,time,status,riskgroup

  result_0 <- data.frame()

  result_0 <- coxph(
    formula = multicox_formulas,
    data = clinical
  )

  result <- summary(result_0)

  result2 <- data.frame(
    var = rownames(result$coefficients),
    coef = result$coefficients[, 1],
    Hazard_Ratio = result$coefficients[, 2],
    se = result$coefficients[, 3],
    p.value = result$coefficients[, 5],
    lower_.95 = result$conf.int[, "lower .95"],
    upper_.95 = result$conf.int[, "upper .95"],
    logrank_pvalue = result$sctest["pvalue"],
    wald_pvalue = result$waldtest["pvalue"],
    Likelihood_pvalue = result$logtest["pvalue"]
  )

  result2 <- result2 %>%
    mutate(HR = paste0(
      sprintf("%.2f",Hazard_Ratio),
      "(",
      sprintf("%.2f",lower_.95),
      "-",
      sprintf("%.2f",upper_.95),
      ")"
    ))
  str(result2)

  temp <- list()
  temp[[1]] <- result2
  temp[[2]] <- result_0

  format_options <- forest_model_format_options(
    colour = points_color,
    color = NULL,
    shape = 18,
    text_size = 3.73,
    point_size = 4.25,
    banded = TRUE
  )
  panels <- list(
      list(width = 0.01),
      list(width = 0.1, display = ~ variable, fontface = "bold", heading = "Variable"),
      list(width = 0.02),
      list(width = 0.15, display = ~ level),
      list(width = 0.1, display = ~n, hjust = 1, heading = "N"),
      # list(width = 0.05, display = ~n_events, width = 0.05, hjust = 1, heading = "Events"),
      # list(
      #   width = 0.05,
      #   display = ~ replace(sprintf("%0.1f", person_time / 365.25), is.na(person_time), ""),
      #   heading = "Person-\nYears", hjust = 1
      # ),
      # Univariable and multivariable Cox regression analysis
      list(width = 0.025, item = "vline", hjust = 0.5),
      list(
        width = 0.25, item = "forest", hjust = 0.5, heading = "Multivariable Cox regression", linetype = "dashed",
        line_x = 0, fontsize = 3.73
      ),
      list(width = 0.025, item = "vline", hjust = 0.5),
      list(width = 0.15, display = ~ ifelse(reference, "Reference", sprintf(
        "%0.2f (%0.2f - %0.2f)",
        trans(estimate), trans(conf.low), trans(conf.high)
      )), display_na = NA, hjust = 0, heading = "HR(95% CI)"),
      list(width = 0.04),
      list(
        width = 0.2,
        display = ~ ifelse(reference, "", case_when(p.value < 0.001 ~ "<0.001",
        p.value > 0.001 ~ sprintf("%.3f",p.value))),
        display_na = NA, hjust = 1, heading = "p"
      ),
      list(width = 0.01)
  )
   
  p <- forest_model(
    model = result_0,
    panels = panels,
    # covariates = vars_for_table,
    merge_models = T,
    return_data = T,
    recalculate_height = T,
    recalculate_width = T,
    format_options = format_options,
    theme = theme_forest()
  )

  temp[[3]] <- p$plot

  #  print(p$plot)

  names(temp) <- c("multicox_clinical_res", "multicox_coxph_res", "plot")
  return(temp)
}


# source("/Pub/Users/wangyk/Project_wangyk/Codelib_YK/some_scr/plotout.R")
